{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import argparse\n",
    "import json\n",
    "from collections import defaultdict, Counter\n",
    "import re\n",
    "\n",
    "from annotation_tool_1 import MAX_WORDS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_repeat_dict(d):\n",
    "    if d[\"loop\"] == \"ntimes\":\n",
    "        repeat_dict = {\"repeat_key\": \"FOR\"}\n",
    "        processed_d = process_dict(with_prefix(d, \"loop.ntimes.\"))\n",
    "        if 'repeat_for' in processed_d:\n",
    "            repeat_dict[\"repeat_count\"] = processed_d[\"repeat_for\"]\n",
    "        if 'repeat_dir' in processed_d:\n",
    "            repeat_dict['repeat_dir'] = processed_d['repeat_dir']\n",
    "        return repeat_dict\n",
    "    if d[\"loop\"] == \"repeat_all\":\n",
    "        repeat_dict = {\"repeat_key\": \"ALL\"}\n",
    "        processed_d = process_dict(with_prefix(d, \"loop.repeat_all.\"))\n",
    "        if 'repeat_dir' in processed_d:\n",
    "            repeat_dict['repeat_dir'] = processed_d['repeat_dir']\n",
    "        return repeat_dict\n",
    "    if d[\"loop\"] == \"forever\":\n",
    "        return {\"stop_condition\": {\"condition_type\": \"NEVER\"}}\n",
    "    if d['loop'] == 'repeat_until':\n",
    "        stripped_d = with_prefix(d, 'loop.repeat_until.')\n",
    "        processed_d = process_dict(stripped_d)\n",
    "        if 'adjacent_to_block_type' in processed_d:\n",
    "            return {\"stop_condition\" : {\n",
    "                        \"condition_type\" : 'ADJACENT_TO_BLOCK_TYPE',\n",
    "                        'block_type': processed_d['adjacent_to_block_type']}\n",
    "                   }\n",
    "        else:\n",
    "            return {\"stop_condition\" : {\n",
    "                        \"condition_type\" : 'ADJACENT_TO_BLOCK_TYPE',}                \n",
    "            }\n",
    "\n",
    "    raise NotImplementedError(\"Bad repeat dict option: {}\".format(d[\"loop\"]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_get_memory_dict(d):\n",
    "    filters_val = d['filters']\n",
    "    out_dict = {'filters': {}}\n",
    "    parent_dict = {}\n",
    "    if filters_val.startswith('type.'):\n",
    "        parts = remove_prefix(filters_val, 'type.').split('.')\n",
    "        type_val = parts[0]\n",
    "        if type_val in ['ACTION', 'AGENT']:\n",
    "            out_dict['filters']['temporal'] = 'CURRENT'\n",
    "            tag_val = parts[1]\n",
    "            out_dict['answer_type'] = 'TAG'\n",
    "            out_dict['tag_name'] = parts[1] # the name of tag is here\n",
    "            if type_val == 'ACTION':\n",
    "                x = with_prefix(d, 'filters.'+filters_val+'.')\n",
    "                out_dict['filters'].update(x)\n",
    "        elif type_val in ['REFERENCE_OBJECT']:\n",
    "            d.pop('filters')\n",
    "            ref_obj_dict = remove_key_prefixes(d, ['filters.type.'])\n",
    "            ref_dict = process_dict(ref_obj_dict)\n",
    "            if 'answer_type' in ref_dict['reference_object']:\n",
    "                out_dict['answer_type'] = ref_dict['reference_object']['answer_type']\n",
    "                ref_dict['reference_object'].pop('answer_type')\n",
    "            if 'tag_name' in ref_dict['reference_object']:\n",
    "                out_dict['tag_name'] = ref_dict['reference_object']['tag_name']\n",
    "                ref_dict['reference_object'].pop('tag_name')    \n",
    "            out_dict['filters'].update(ref_dict)\n",
    "            \n",
    "        out_dict['filters']['type'] = type_val\n",
    "        \n",
    "    return out_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_prefix(text, prefix):\n",
    "    if text.startswith(prefix):\n",
    "        return text[len(prefix):]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [],
   "source": [
    "def handle_get_memory(d):\n",
    "    out_d = {'dialogue_type': 'GET_MEMORY'}\n",
    "    child_d = process_get_memory_dict(with_prefix(d, \"action_type.ANSWER.\"))\n",
    "    out_d.update(child_d)\n",
    "    return out_d\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert s to snake case\n",
    "def snake_case(s):\n",
    "    return re.sub(\"([a-z])([A-Z])\", \"\\\\1_\\\\2\", s).lower()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''this function splits the key that starts with a given prefix and only for values that are not None\n",
    "and makes the key be the thing after prefix\n",
    "'''\n",
    "def with_prefix(d, prefix):\n",
    "    return {\n",
    "        k.split(prefix)[1]: v\n",
    "        for k, v in d.items()\n",
    "        if k.startswith(prefix) and v not in (\"\", None, \"None\")\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {},
   "outputs": [],
   "source": [
    "''' this function removes certain prefixes from keys and renames the key to be: key with text following \n",
    "the prefix in the dict'''\n",
    "def remove_key_prefixes(d, ps):\n",
    "    \n",
    "    for p in ps:\n",
    "        d = d.copy()\n",
    "        rm_keys = []\n",
    "        add_items = []\n",
    "        # print(p, d)\n",
    "        for k, v in d.items():\n",
    "            if k.startswith(p):\n",
    "                rm_keys.append(k)\n",
    "                add_items.append((k[len(p) :], v))\n",
    "        for k in rm_keys:\n",
    "            del d[k]\n",
    "        for k, v in add_items:\n",
    "            d[k] = v\n",
    "    return d\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_spans_due_to_empty_words(action_dict, words):\n",
    "    \"\"\"Return modified (action_dict, words)\"\"\"\n",
    "\n",
    "    def reduce_span_vals_gte(d, i):\n",
    "        for k, v in d.items():\n",
    "            if type(v) == dict:\n",
    "                reduce_span_vals_gte(v, i)\n",
    "                continue\n",
    "            try:\n",
    "                a, b = v\n",
    "                if a >= i:\n",
    "                    a -= 1\n",
    "                if b >= i:\n",
    "                    b -= 1\n",
    "                d[k] = [[a, b]]\n",
    "            except ValueError:\n",
    "                pass\n",
    "            except TypeError:\n",
    "                pass\n",
    "\n",
    "    # remove trailing empty strings\n",
    "    while words[-1] == \"\":\n",
    "        del words[-1]\n",
    "\n",
    "    # fix span\n",
    "    i = 0\n",
    "    while i < len(words):\n",
    "        if words[i] == \"\":\n",
    "            reduce_span_vals_gte(action_dict, i)\n",
    "            del words[i]\n",
    "        else:\n",
    "            i += 1\n",
    "\n",
    "    return action_dict, words\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_dict(d):\n",
    "    r = {}\n",
    "\n",
    "#     print(d)\n",
    "#     print(\"----------------\")\n",
    "    d = remove_key_prefixes(d, [\"COPY.yes.\", \"COPY.no.\", 'FREEBUILD.BUILD.', 'answer_type.TAG.', 'FREEBUILD.FREEBUILD.', 'coref_resolve_check.yes.', 'coref_resolve_check.no.'])\n",
    "#     print(d)\n",
    "#     print(\"----------------new------------------\")\n",
    "    if \"location\" in d:\n",
    "        r[\"location\"] = {\"location_type\": d[\"location\"]}\n",
    "        if r['location']['location_type'] == 'coref_resolve_check':\n",
    "            del r['location']['location_type']\n",
    "        elif r[\"location\"][\"location_type\"] == \"REFERENCE_OBJECT\":\n",
    "            r[\"location\"][\"location_type\"] = \"REFERENCE_OBJECT\"\n",
    "            r[\"location\"][\"relative_direction\"] = d.get(\n",
    "                \"location.REFERENCE_OBJECT.relative_direction\"\n",
    "            )\n",
    "            # no key for EXACT\n",
    "            if r[\"location\"][\"relative_direction\"] in (\"EXACT\", \"Other\"):\n",
    "                del r[\"location\"][\"relative_direction\"]\n",
    "            d[\"location.REFERENCE_OBJECT.relative_direction\"] = None\n",
    "        r[\"location\"].update(process_dict(with_prefix(d, \"location.\")))\n",
    "        \n",
    "    for k, v in d.items():\n",
    "        if (\n",
    "            k == \"location\"\n",
    "            or k in ['COPY', 'coref_resolve_check']\n",
    "            or (k == \"relative_direction\" and v in (\"EXACT\", \"NEAR\", \"Other\"))\n",
    "        ):\n",
    "            continue\n",
    "        # handle span\n",
    "        if re.match(\"[^.]+.span#[0-9]+\", k):\n",
    "            prefix, rest = k.split(\".\", 1)\n",
    "            idx = int(rest.split(\"#\")[-1])\n",
    "            if prefix in r:\n",
    "                r[prefix].append([idx, idx])\n",
    "#                 a, b = r[prefix]\n",
    "#                 r[prefix] = [min(a, idx), max(b, idx)]  # expand span to include idx\n",
    "            else:\n",
    "                r[prefix] = [[idx, idx]]\n",
    "\n",
    "        # handle nested dict\n",
    "        elif \".\" in k:\n",
    "            prefix, rest = k.split(\".\", 1)\n",
    "            prefix_snake = snake_case(prefix)\n",
    "            r[prefix_snake] = r.get(prefix_snake, {})\n",
    "            r[prefix_snake].update(process_dict(with_prefix(d, prefix + \".\")))\n",
    "\n",
    "        # handle const value\n",
    "        else:\n",
    "            r[k] = v\n",
    "\n",
    "    return r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "metadata": {},
   "outputs": [],
   "source": [
    "def handle_put_memory(d):\n",
    "    return {}\n",
    "    \n",
    "\n",
    "def handle_commands(d):\n",
    "    output = {}\n",
    "    action_name = d[\"action_type\"]\n",
    "    formatted_dict = with_prefix(d, \"action_type.{}.\".format(action_name))\n",
    "    child_d = process_dict(with_prefix(d, \"action_type.{}.\".format(action_name)))\n",
    "    # Fix Build/Freebuild mismatch\n",
    "    if child_d.get(\"FREEBUILD\") == \"FREEBUILD\":\n",
    "        action_name = 'FREEBUILD'\n",
    "    child_d.pop(\"FREEBUILD\", None)\n",
    "\n",
    "    if formatted_dict.get('COPY', 'no') == 'yes':\n",
    "        action_name = 'COPY'\n",
    "        formatted_dict.pop('COPY')\n",
    "    \n",
    "    # add action type info\n",
    "    output['action_type'] = ['yes', action_name.lower()]\n",
    "    # add dialogue type info\n",
    "    if output['action_type'][1] == 'tag':\n",
    "        output['dialogue_type'] = ['yes', 'PUT_MEMORY']\n",
    "    else:\n",
    "        output['dialogue_type'] = ['yes', 'HUMAN_GIVE_COMMAND']\n",
    "    \n",
    "    for k, v in child_d.items():\n",
    "        if k =='target_action_type':\n",
    "            output[k] = ['yes', v]\n",
    "        elif type(v)==list:\n",
    "            output[k]= ['no', v]\n",
    "        else:\n",
    "            output[k] = ['yes', v]\n",
    "    return output\n",
    "\n",
    "def process_result(full_d, index):\n",
    "    worker_id = full_d[\"WorkerId\"]\n",
    "    d = with_prefix(full_d, \"Answer.root.{}.\".format(index))\n",
    "    if not d:\n",
    "        return worker_id, {}, full_d['Input.command_{}'.format(index)].split()\n",
    "    try:\n",
    "        action = d[\"action_type\"]\n",
    "    except KeyError:\n",
    "        return worker_id, {}, full_d['Input.command_{}'.format(index)].split()\n",
    "\n",
    "    action_dict = handle_commands(d)\n",
    "        \n",
    "    ##############\n",
    "    # repeat dict\n",
    "    ##############\n",
    "    #NOTE: this can probably loop over or hold indices of which specific action ?\n",
    "    if action_dict.get('dialogue_type', [None, None])[1] == 'HUMAN_GIVE_COMMAND':\n",
    "        if d.get(\"loop\") not in [None, \"Other\"]:\n",
    "            repeat_dict = process_repeat_dict(d)\n",
    "            # Some turkers annotate a repeat dict for a repeat_count of 1.\n",
    "            # Don't include the repeat dict if that's the case\n",
    "            if repeat_dict.get('repeat_dir', None) == 'Other':\n",
    "                repeat_dict.pop('repeat_dir')\n",
    "            if repeat_dict.get(\"repeat_count\"):\n",
    "                a, b = repeat_dict[\"repeat_count\"][0]\n",
    "                repeat_count_str = \" \".join(\n",
    "                    [full_d[\"Input.word{}{}\".format(index, x)] for x in range(a, b + 1)]\n",
    "                )\n",
    "                if repeat_count_str not in (\"a\", \"an\", \"one\", \"1\"):\n",
    "                    action_dict['repeat'] = ['yes', repeat_dict]\n",
    "#                     action_val = list(action_dict.values())[0]  # check what this is\n",
    "#                     if action_specific_dict.get(\"schematic\"):\n",
    "#                         action_specific_dict[\"schematic\"][\"repeat\"] = repeat_dict\n",
    "#                     elif action_specific_dict.get(\"reference_object\"):\n",
    "#                         action_specific_dict[\"reference_object\"][\"repeat\"] = repeat_dict\n",
    "#                     else:\n",
    "#                         action_specific_dict[\"repeat\"] = repeat_dict\n",
    "            else:\n",
    "                action_dict['repeat'] = ['yes', repeat_dict]\n",
    " \n",
    "            \n",
    "    ##################\n",
    "    # post-processing\n",
    "    ##################\n",
    "\n",
    "   \n",
    "\n",
    "    # Fix empty words messing up spans\n",
    "    words = [full_d[\"Input.word{}{}\".format(index, x)] for x in range(MAX_WORDS)]\n",
    "    action_dict, words = fix_spans_due_to_empty_words(action_dict, words)\n",
    "\n",
    "    return worker_id, action_dict, words\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_cnt_in_schematic(words, action_dict):\n",
    "    if 'repeat' not in action_dict:\n",
    "        return action_dict\n",
    "    repeat = action_dict['repeat']\n",
    "    val = []\n",
    "    if 'repeat_count' in repeat[1]:\n",
    "        val = repeat[1]['repeat_count']\n",
    "    elif 'repeat_key' in repeat[1] and repeat[1]['repeat_key'] == 'ALL':\n",
    "        if any(x in ['all', 'every', 'each'] for x in words):\n",
    "            all_val = words.index('all')\n",
    "            val = [[all_val, all_val]]\n",
    "    else:\n",
    "        return action_dict\n",
    "    \n",
    "    for k, v in action_dict.items():\n",
    "        if k in ['schematic', 'reference_object']:\n",
    "            for i, meh in enumerate(v[1]):\n",
    "                # print(words, val)\n",
    "                if meh in val:\n",
    "                    v[1].pop(i)\n",
    "            action_dict[k] = [v[0], v[1]]\n",
    "    return action_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 207,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "unique_keys = []\n",
    "with open('/Users/kavyasrinet/Downloads/test_q.csv', \"r\") as f, open('/Users/kavyasrinet/Downloads/test_q.txt', 'w') as f2:\n",
    "    r = csv.DictReader(f)\n",
    "    all_data = {}\n",
    "    for d in r:\n",
    "        worker_id = d[\"WorkerId\"]\n",
    "        all_data[worker_id] = {}\n",
    "        for i in range(1, 4):\n",
    "            sentence = d['Input.command_{}'.format(i)]\n",
    "            _, action_dict, words = process_result(d, i)\n",
    "            a_dict = fix_cnt_in_schematic(words, action_dict)\n",
    "            unique_keys.extend(list(a_dict.keys()))\n",
    "            all_data[worker_id][sentence] = a_dict\n",
    "\n",
    "    for k, v in all_data.items():\n",
    "        f2.write(k+\"\\t\"+str(v)+\"\\n\")\n",
    "    print(len(all_data.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'repeat', 'schematic', 'dialogue_type', 'action_type', 'reference_object', 'location'}\n"
     ]
    }
   ],
   "source": [
    "a = set(unique_keys)\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 500 qual test: '/Users/kavyasrinet/Downloads/500_qual_test.csv'\n",
    "# first round: '//Users/kavyasrinet/Downloads/14_qual_test.csv'\n",
    "# test from sandbox: 'data/test.csv'"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
